from .image_convert_base import ConvertBase
import cv2 as cv
import numpy as np
import random


class NoiseConvert(ConvertBase):
    
    def __init__(self, use_rate=0.5, nose_types=['高斯噪声', '椒盐噪声', '滤波模糊'], is_random = True):
        super().__init__(use_rate)
        self._converts = {
            '高斯噪声' : self.gasuss,
            '椒盐噪声' : self.rand_nose,
            '滤波模糊' : self.conv_low
        }
        
        self.is_random = is_random
        self.nose_types = nose_types
    
    def gasuss(self, image, mean=0, var=0.001):
        '''
        添加高斯噪声
        mean : 均值
        var : 方差
        '''
        image = np.array(image/255, dtype=float)
        noise = np.random.normal(mean, var ** 0.5, image.shape)
        out = image + noise
        if out.min() < 0:
            low_clip = -1.
        else:
            low_clip = 0.
        out = np.clip(out, low_clip, 1.0)
        out = np.uint8(out*255)
        return out

    def rand_nose(self, image):
        output = np.zeros(image.shape,np.uint8)
        prob=random.randint(5, 10) / 10000.0  #随机噪声比例
        thres = 1 - prob 
        for i in range(image.shape[0]):
            for j in range(image.shape[1]):
                rdn = random.random()
                if rdn < prob:
                    output[i][j] = 0
                elif rdn > thres:
                    output[i][j] = 255
                else:
                    output[i][j] = image[i][j]
        return output
    
    def conv_low(self, img, size = 5):
        #利用opencv模块
        r = random.randint(1, 3)
        if r == 1:
            img = cv.blur(img,(size,size))  #中值滤波
        elif r == 2:
            img = cv.medianBlur(img,size)  #椒盐滤波
        elif r == 3:
            img = cv.GaussianBlur(img,(size,size), 0)  #高斯滤波

        return img
    
    def convert(self, img, boxes, points):
        r = random.randint(0, len(self.nose_types) - 1)
        type_md = self._converts[self.nose_types[r]]
        
        img = type_md(img)
        return img, boxes, points
        
        
        
        
        
